package zagaynov;

import common.DerivnFunction;
import common.RungeKutta;

import java.util.*;

public class Zagaynov {

    // Variable declarations

    double initialTimeValue;                     // initial integration value (x)
    double xn;                     // final integration value (x)

    double[] y0;         // initial values of the y[i]
    double[] yn;         // returned value of the y[i] at x = xn

    public Map<Double, List<Double>> doCalculations(Properties properties) {

        double[][] koef = calculateK_Matrix(properties);
        Equity dn = new Equity(properties, koef);

        // Assign initial values of y[i]
        setInitialValues(properties);

        Map<Double, List<Double>> results = new TreeMap<Double, List<Double>> ();
        
        for (double i = 0; i < properties.getFinalTimeValue(); i += properties.getTimeStepSize()) {
            doIteration(dn, properties, results);
            initialTimeValue = xn;
            xn = initialTimeValue + properties.getTimeStepSize();
            y0 = yn;
        }

         return results;
    }

    private void doIteration(Equity dn, Properties properties, Map<Double, List<Double>> results) {
        // Create and instance of common.RungeKutta
        RungeKutta rk = new RungeKutta();

        // Set values needed by fixed step size method
        rk.setInitialValueOfX(initialTimeValue);
        rk.setFinalValueOfX(xn);
        rk.setInitialValuesOfY(y0);
        rk.setStepSize(properties.getIntegrationStepSize());

        // Call Fourth Order Runge-Kutta method
        yn = rk.fourthOrder(dn);
        
        List<Double> result = new ArrayList<Double>();

        for (int i=0; i<yn.length; i++) {
          result.add(yn[i]);
        }
        
        results.put(xn, result);
        //logResults(xn, yn, rk);
    }


    private double[][] calculateK_Matrix(Properties properties) {
        CoagulationFactor coagulationFactor = properties.getCoagulationFactor();
        int n = properties.getEquationsNumber();
        double[][] koef = new double[n][n];
        for(int k=0; k < n; k++) {
            for(int j=0; j < n; j++) {
                double normK = k+1;
                double normJ = j+1;
                switch (coagulationFactor) {
                    case UNIT:
                        koef[k][j]=1;
                        continue;
                    case SUM:
                        koef[k][j]=normK + normJ;
                        continue;
                    case PRODUCT:
                        koef[k][j]=normK * normJ;
                        continue;
                    case FREEMOLECULAR:
                        koef[k][j]=Math.pow((Math.pow(normK, 1.0/3)+Math.pow(normJ, 1.0/3)),2)*
                                Math.pow((Math.pow(normK, -1)+Math.pow(normJ, -1)),1.0/2);
                        continue;
                    case BROUNCOAGULATION:
                        koef[k][j]=(Math.pow(normK, 1.0/3)+Math.pow(normJ, 1.0/3))*
                                (Math.pow(normK, -1.0/3)+Math.pow(normJ, -1.0/3));
                        continue;
                    case GRADIENTCOAGULATION:
                        koef[k][j]=Math.pow((Math.pow(normK, 1.0/3)+Math.pow(normJ, 1.0/3)),3);
                        continue;
                    case GRAVITATIONCOAGULATION:
                        koef[k][j]=Math.pow((Math.pow(normK, 1.0/3)+Math.pow(normJ, 1.0/3)),2)*
                                Math.abs(-(Math.pow(normK, 2.0/3)+Math.pow(normJ, 2.0/3)));
                        continue;
                    case RAREGRAVITAIONCOAGULATION:
                        koef[k][j]=Math.pow(normK, 2.0/3)*Math.pow(normJ, 2.0/3);
                        continue;
                    default:
                        throw new RuntimeException("Invalid coagulation Factor");
                }
            }
        }
        return koef;

    }

    private void setInitialValues(Properties properties) {
        int n = properties.getEquationsNumber();
        y0 = new double[n];         // initial values of the y[i]
        yn = new double[n];         // returned value of the y[i] at x = xn

        for (int i=0; i < n; i++) {
            y0[i] = 0.0D;
        }
        initialTimeValue = properties.getInitialTimeValue();
        xn = initialTimeValue + properties.getTimeStepSize();
    }

    private void logResults(double xn, double[] yn, RungeKutta rk) {
        // Output the results
        //System.out.println("Fourth order Runge-Kutta procedure");
        for (int i = 0; i < yn.length; i++) {
            System.out.println("The value of y[" + i + "] at x = " + xn + " is " + yn[i]);
        }
        //System.out.println("Number of iterations = " + rk.getNumberOfIterations());
        System.out.println("==========");
    }

    private void logZagaynov(double xn, double[] yn, RungeKutta rk) {
        // Output the results
        //System.out.println("Fourth order Runge-Kutta procedure");

        System.out.println("The value of y[" + 0 + "] at x = " + xn + " is " + yn[0]);
        System.out.println("The value of y[" + 1 + "] at x = " + xn + " is " + yn[1]);
        System.out.println("The value of y[" + 2 + "] at x = " + xn + " is " + yn[2]);
        System.out.println("The value of y[" + 3 + "] at x = " + xn + " is " + yn[3]);
        System.out.println("The value of y[" + 49 + "] at x = " + xn + " is " + yn[49]);
        System.out.println("The value of y[" + 74 + "] at x = " + xn + " is " + yn[74]);
        System.out.println("The value of y[" + 99 + "] at x = " + xn + " is " + yn[99]);

        //System.out.println("Number of iterations = " + rk.getNumberOfIterations());
        System.out.println("==========");
    }

    class Equity implements DerivnFunction {
        private Properties properties;
        private double[][] koef;

        public Equity(Properties properties, double[][] koef) {
            this.properties = properties;
            this.koef = koef;
        }

        public double[] derivn(double x, double[] y) {
            if (properties.getModelType().equals(ModelType.COAGULATION)) {
                return complex(y);
            } else {
                return simple(y);
            }

        }

        private double[] complex(double[] y) {
            int n = properties.getEquationsNumber();
            double[] dydx = new double[y.length];
            dydx[0] = properties.getSourcePower();

            for (int i=0; i<n; i++) {
                dydx[0]=dydx[0]-koef[i][0]*y[i]*y[0];
            }

            for (int i=1;i<n; i++) {
                double firstPart=0;

                for (int j=0; j<i; j++) {
                    firstPart = firstPart + koef[j][i-j-1] * y[j]*y[i-j-1];  //TODO check this
                }

                double secondPart=0;

                for (int j = 0; j < i; j++) {
                    secondPart = secondPart + koef[i][j] * y[i]*y[j];  //TODO check this
                }


                dydx[i]=0.5*(firstPart) - secondPart;
            }


            return dydx;
        }

        private double[] simple(double[] y) {
            double[] dydx = new double[y.length];

            //dydx[0] = ic - 2*y[0]*y[0] -y[0]*y[1] - y[0]*y[2];
            //dydx[1] = y[0]*y[0]-y[1]*y[0];
            //dydx[2] = y[1]*y[0] - y[2]*y[0];
            int n = properties.getEquationsNumber();

            dydx[0] = properties.getSourcePower() - koef[0][0]*y[0]*y[0];

            for (int i=0; i<n; i++) {
                dydx[0]=dydx[0]-koef[i][0]*y[i]*y[0];
            }

            for (int i=1; i<n; i++) {
                dydx[i] = koef[i-1][0]*y[i-1]*y[0]-koef[i][0]*y[i]*y[0];
            }

            return dydx;
        }

    }
}

